// STL
#include <iostream>
// BOOST
#include <boost/mpi.hpp>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>
// ARMADILLO
#include <armadillo>
// FACTORIZATION
#include "helper/helper.h"
#include "helper/mpi_helper.h"
#include "helper/randGen.h"
#include "factor_jobs.h"

#include "factorization/multNMF.h"
#include "factorization/HALS.h"
#include "factorization/stochHALS.h"

using namespace std;
using namespace arma;
namespace mpi = boost::mpi;
namespace po = boost::program_options;

void JOB::processParameter(int argc, char* argv[]) {
  desc.add_options()
    ("help", "produce help message")
    ("job", po::value<string>(), "job type:\n"
     "  runtest: run trivial tests\n"
     "  multNMF: single thread multiplicative NMF\n"
     "  parMultNMF: multiple threads multiplicative NMF\n"
     "  hals: Hierachical Alternative Least Square\n"
     "  parHals: parallel Hierachical Alternative Least Square\n"
     "  parStochHals: parallel Hierachical Alternative Least Square\n"
     )
    ("type", po::value<string>()->default_value("dense"), "\"dense\" or \"sparse\" matrix")
    // data
    ("pr_input", po::value<string>(), "prefix of input files, use \"rand\" for random input")
    // nmf parameter
    ("m", po::value<int>(), "m = A.n_rows")
    ("n", po::value<int>(), "n = A.n_cols")
    ("randk", po::value<int>(), "k = W.n_cols = H.n_cols (low rank setting)")
    ("k", po::value<int>(), "k = W.n_cols = H.n_cols (low rank setting)")
    ("s", po::value<int>()->default_value(100), "random seed for data generation")
    ("cs", po::value<int>()->default_value(200), "random seed for W, H initialization")
    // solver parameter
    ("iter", po::value<int>()->default_value(10), "maximum iterations")
    ("epsilon", po::value<double>()->default_value(1e-8), "additive term to prevent division by 0")
    ("tol", po::value<double>()->default_value(1e-4), "stop criteria tolerance")
    ("sparse", po::value<double>(), "sparsity level")
    ("alpha", po::value<double>()->default_value(1.0), "power order of step size (1/t)^alpha")
    // display parameter
    ("verbose", po::value<int>()->default_value(0), "verbose setting: 0,1,...")
    ("outputDir", po::value<string>()->default_value("../../result/cpp"), "output file")
    ("output", po::value<string>()->default_value("default_output"), "output file")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}

void JOB::readInput(mat& Arows, mat& Acols)
{
  REQUIRE_PARAM("pr_input");
  string prefix = vm["pr_input"].as<string>();
  RAND<mat> ra(world, vm);
  
  if (prefix == "rand") {
    ra.randNMF(Arows, Acols);
  }
  else {
    stringstream sr, sc;
    sr << prefix << "rows_" << SIZE(world);
    sc << prefix << "cols_" << SIZE(world);
    
    PART<mat>::load(world, sr.str(), Arows, raw_ascii);
    PART<mat>::load(world, sc.str(), Acols, raw_ascii);
  }
}

void JOB::readInput(mat& A)
{
  REQUIRE_PARAM("pr_input");
  string prefix = vm["pr_input"].as<string>();
  RAND<mat> ra(world, vm);
  
  if (prefix == "rand") {
    ra.randNMF(A);
  }
  else {
    A.load(prefix, raw_ascii);
  }
}

void JOB::doMultNMF()
{
  if (vm.count("sparse")) {
    factor::MultNMF<sp_mat> solver(world, vm);
    double sparsity = vm["sparse"].as<double>();
    sp_mat W0 = sprandu<sp_mat>(100,20,sparsity), H0 = sprandu<sp_mat>(90,20,sparsity);
    sp_mat A = W0*trans(H0);
    solver.NMF(A);
  }
  else {
    factor::MultNMF<mat> solver(world, vm);
    mat W0 = randu<mat>(10,2), H0 = randu<mat>(9,2), W, H;
    mat A = W0*trans(H0);
    solver.NMF(A);
  }
}

void JOB::doParMultNMF()
{
  mat Arows, Acols;
  readInput(Arows, Acols);
    
  factor::MultNMF<mat> solver(world, vm);
  solver.NMF(Arows, Acols);
}

void JOB::doHALS()
{
  mat A;
  readInput(A);

  factor::HALS<mat> solver(world, vm);
  solver.NMF(A);
}

void JOB::doParHALS()
{
  mat Arows, Acols;
  readInput(Arows, Acols);

  factor::HALS<mat> solver(world, vm);
  solver.NMF(Arows, Acols);
}

void JOB::doStochHALS()
{
  mat A;
  readInput(A);

  factor::StochHALS<mat> solver(world, vm);
  solver.NMF(A);
}

void JOB::doParStochHALS()
{
  if (vm["type"].as<string>() == "dense") {
    mat Arows, Acols;
    readInput(Arows, Acols);
    
    factor::StochHALS<mat> solver(world, vm);
    solver.NMF(Arows, Acols);
  }
  else {
    
  }
}
